# -*- encoding: utf-8 -*-
from PyQt5.QtWidgets import QApplication, QMainWindow, QMessageBox, QTableWidgetItem, QFileDialog
from PyQt5.QtCore import Qt
from Ui.Ui import Ui_MainWindow
from pywifi import PyWiFi
from utils.threads import BreakWiFiKeyThread
import sys
import os


class MainWindow(QMainWindow):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.initUI()

    def initUI(self):
        self.ui = Ui_MainWindow()
        self.ui.setupUi(self)

        # 设置表头列宽
        self.ui.tableWidget.horizontalHeader().resizeSection(0, 110)
        self.ui.tableWidget.horizontalHeader().resizeSection(1, 200)
        self.ui.tableWidget.horizontalHeader().resizeSection(2, 80)

        # 扫描无线网卡
        self.wireless_card_mapping = {}
        self.event_scanWirelessCard()
        self.ui.btn_flushWirelessCard.clicked.connect(self.event_scanWirelessCard)

        # 打开文件
        self.ui.btn_openFolder.clicked.connect(self.event_openFolder)

        # 扫描WiFi
        self.ui.btn_scanWiFi.clicked.connect(self.event_scanWiFi)

        # 选择SSID
        self.ui.tableWidget.itemSelectionChanged.connect(self.event_selectionSSID)

        # 破解WiFi
        self.ui.btn_breakWiFi.clicked.connect(self.event_breakWiFi)
        # 停止破解
        self.ui.btn_stopBreak.clicked.connect(self.event_stopBreak)

    def event_scanWirelessCard(self):
        wifi_obj = PyWiFi()
        ifaceList = wifi_obj.interfaces()
        for iface in ifaceList:
            if iface.name() not in self.wireless_card_mapping:
                self.ui.comboBox.addItem(iface.name())
                self.wireless_card_mapping[iface.name()] = iface

    def event_openFolder(self):
        filename = QFileDialog.getOpenFileName(self, "选择文件", "./", "FILE(*.txt)")
        if filename[0]:
            with open(filename[0], 'rt', encoding='utf-8') as f:
                content = f.read()
            if len(content) > 0:
                self.ui.edit_filePath.setText(filename[0])
            else:
                QMessageBox.warning(self, "警告", "文件内容为空")

    def event_scanWiFi(self):
        iface = self.wireless_card_mapping.get(self.ui.comboBox.currentText())
        if iface:
            iface.scan()
            results = iface.scan_results()

            if self.ui.tableWidget.rowCount() > 0 and self.ui.tableWidget.columnCount() > 0:
                self.ui.tableWidget.clearContents()
                self.ui.tableWidget.setRowCount(0)

            wifi_mapping = {}
            for i in results:
                wifi_mapping[i.ssid] = i

            self.ui.tableWidget.setRowCount(len(wifi_mapping))
            self.ui.tableWidget.setColumnCount(3)

            for row, profile in enumerate(wifi_mapping.values()):
                ssid_item = QTableWidgetItem(str(profile.ssid))
                ssid_item.setTextAlignment(Qt.AlignCenter)
                ssid_item.setFlags(Qt.ItemIsSelectable | Qt.ItemIsEnabled)
                self.ui.tableWidget.setItem(row, 0, ssid_item)

                bssid_item = QTableWidgetItem(str(profile.bssid))
                bssid_item.setTextAlignment(Qt.AlignCenter)
                bssid_item.setFlags(Qt.ItemIsSelectable | Qt.ItemIsEnabled)
                self.ui.tableWidget.setItem(row, 1, bssid_item)

                signal_item = QTableWidgetItem(str(profile.signal))
                signal_item.setTextAlignment(Qt.AlignCenter)
                signal_item.setFlags(Qt.ItemIsSelectable | Qt.ItemIsEnabled)
                self.ui.tableWidget.setItem(row, 2, signal_item)
        else:
            QMessageBox.warning(self, "提示", "请选择无线网卡或刷新无线网卡")

    def event_selectionSSID(self):
        current_item = self.ui.tableWidget.currentItem()
        row_item = current_item.row()
        ssid = self.ui.tableWidget.item(row_item, 0).text()
        self.ui.edit_target.setText(ssid)

    def event_breakWiFi(self):
        if not self.ui.edit_target.text():
            QMessageBox.warning(self, "提示", "未选中SSID")
            return
        if not self.ui.edit_filePath.text():
            QMessageBox.warning(self, "错误", "未选择密码文件")
            return
        if not os.path.exists(self.ui.edit_filePath.text()):
            QMessageBox.warning(self, "错误", "文件不存在")
            return

        self.breaking = BreakWiFiKeyThread(
            iface=self.wireless_card_mapping.get(self.ui.comboBox.currentText()),
            ssid=self.ui.edit_target.text().strip(),
            file_path=self.ui.edit_filePath.text().strip()
        )
        self.breaking.result_signal.connect(self.event_updateResults)
        self.breaking.start()

    def event_stopBreak(self):
        self.breaking.terminate()

    def event_updateResults(self, msg):
        self.ui.textBrowser.append(msg)


if __name__ == '__main__':
    app = QApplication(sys.argv)

    window = MainWindow()
    window.show()

    sys.exit(app.exec_())
